package com.xinghen.common.zookeeper.registry.base;

import com.google.common.base.Charsets;
import com.google.common.collect.Maps;
import com.xinghen.common.base.GlobalConstant;
import com.xinghen.common.base.properties.ZookeeperProperties;
import com.xinghen.common.zookeeper.registry.exception.RegException;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.api.ACLProvider;
import org.apache.curator.framework.recipes.atomic.AtomicValue;
import org.apache.curator.framework.recipes.atomic.DistributedAtomicInteger;
import org.apache.curator.framework.recipes.cache.ChildData;
import org.apache.curator.framework.recipes.cache.TreeCache;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.curator.retry.RetryNTimes;
import org.apache.curator.utils.CloseableUtils;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.ZooDefs;
import org.apache.zookeeper.data.ACL;
import org.apache.zookeeper.data.Stat;

import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.TimeUnit;

/**
 * @author xinghen
 * @ClassName: RegistryCenter
 * @Package: com.xinghen.common.zookeeper.registry.base
 * @Description 基于Zookeeper的注册中心
 * @create 2018-04-12 14:51
 */
@Slf4j
public class RegistryCenter {

    /**
     * zookeeper配置
     */
    private ZookeeperProperties zookeeperProperties;

    /**
     * zookeeper客户端
     */
    private CuratorFramework client;

    /**
     * 分布式原子变量
     */
    private DistributedAtomicInteger distributedAtomicInteger;

    /**
     * 缓存
     */
    private Map<String, TreeCache> caches = Maps.newConcurrentMap();

    public RegistryCenter(ZookeeperProperties zookeeperProperties) {
        this.zookeeperProperties = zookeeperProperties;
    }

    /**
     * 初始化注册中心
     */
    public void init() {
        log.debug("zookeeper registry center init, server lists is: {}.", zookeeperProperties.getZkAddressList());
        CuratorFrameworkFactory.Builder builder = CuratorFrameworkFactory.builder()
                .connectString(zookeeperProperties.getZkAddressList())
                .retryPolicy(new ExponentialBackoffRetry(zookeeperProperties.getBaseSleepTimeMilliseconds(),
                        zookeeperProperties.getMaxRetries(), zookeeperProperties.getMaxSleepTimeMilliseconds()));
        if (zookeeperProperties.getSessionTimeoutMilliseconds() > 0) {
            builder.sessionTimeoutMs(zookeeperProperties.getSessionTimeoutMilliseconds());
        }
        if (zookeeperProperties.getConnectionTimeoutMilliseconds() > 0) {
            builder.connectionTimeoutMs(zookeeperProperties.getConnectionTimeoutMilliseconds());
        }
        if (StringUtils.isNotBlank(zookeeperProperties.getDigest())) {
            builder.authorization("digest", zookeeperProperties.getDigest().getBytes(Charsets.UTF_8))
                    .aclProvider(new ACLProvider() {
                        @Override
                        public List<ACL> getDefaultAcl() {
                            return ZooDefs.Ids.CREATOR_ALL_ACL;
                        }

                        @Override
                        public List<ACL> getAclForPath(String s) {
                            return ZooDefs.Ids.CREATOR_ALL_ACL;
                        }
                    });
        }
        client = builder.build();
        client.start();
        try {
            if (!client.blockUntilConnected(zookeeperProperties.getMaxSleepTimeMilliseconds() *
                    zookeeperProperties.getMaxRetries(), TimeUnit.MILLISECONDS)) {
                CloseableUtils.closeQuietly(client);
                throw new KeeperException.OperationTimeoutException();
            }
        } catch (Exception e) {
            log.error(e.getMessage());
        }
    }

    /**
     * 销毁注册中心
     */
    public void close() {
        for (TreeCache treeCache : caches.values()) {
            treeCache.close();
        }
        //等待500ms, cache先关闭再关闭client
        try {
            TimeUnit.MILLISECONDS.sleep(500);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }
        CloseableUtils.closeQuietly(client);
    }

    /**
     * 获取数据
     *
     * @param key
     * @return
     */
    public String get(String key) {
        TreeCache cache = findTreeCache(key);
        if (Objects.nonNull(cache)) {
            ChildData childData = cache.getCurrentData(key);
            if (Objects.nonNull(childData)) {
                byte[] data = childData.getData();
                return Objects.nonNull(data) ? new String(data, Charsets.UTF_8) : null;
            }
        }
        return getDirectly(key);
    }

    /**
     * 直接从注册中心获取数据
     *
     * @param key
     * @return
     */
    public String getDirectly(String key) {
        try {
            byte[] data = getClient().getData().forPath(key);
            if (Objects.nonNull(data)) {
                return new String(data, Charsets.UTF_8);
            }
        } catch (Exception e) {
            log.error(e.getMessage());
        }
        return null;
    }

    /**
     * 获取子节点
     *
     * @param key
     * @return
     */
    public List<String> getChildrenKey(String key) {
        try {
            List<String> list = getClient().getChildren().forPath(key);
            if (CollectionUtils.isNotEmpty(list)) {
                list.sort(Comparator.reverseOrder());
                return list;
            }
        } catch (Exception e) {
            log.error(e.getMessage());
        }
        return null;
    }

    /**
     * 获取子节点数
     *
     * @param key
     * @return
     */
    public int getNumChildren(final String key) {
        try {
            Stat stat = getClient().checkExists().forPath(key);
            if (Objects.nonNull(stat)) {
                return stat.getNumChildren();
            }
        } catch (Exception e) {
            log.error(e.getMessage());
        }
        return 0;
    }

    /**
     * 判断节点是否存在
     *
     * @param key
     * @return
     */
    public boolean exists(String key) {
        try {
            Stat stat = getClient().checkExists().forPath(key);
            return Objects.nonNull(stat);
        } catch (Exception e) {
            log.error(e.getMessage());
        }
        return false;
    }

    /**
     * 持久化节点数据
     *
     * @param key
     * @param value
     */
    public void persist(String key, String value) {
        try {
            if (exists(key)) {
                update(key, value);
            } else {
                create(key, value);
            }
        } catch (Exception e) {
            log.error(e.getMessage());
        }
    }

    /**
     * 持久化节点数据
     *
     * @param key
     */
    public void persist(String key) {
        this.persist(key, null);
    }

    /**
     * 创建持久节点
     *
     * @param key
     * @param value
     */
    public void create(String key, String value) {
        if (exists(key)) {
            log.error("{}节点已存在", key);
            return;
        }
        try {
            if (StringUtils.isBlank(value)) {
                getClient().create().creatingParentsIfNeeded().withMode(CreateMode.PERSISTENT).forPath(key);
            } else {
                getClient().create().creatingParentsIfNeeded().withMode(CreateMode.PERSISTENT)
                        .forPath(key, value.getBytes(Charsets.UTF_8));
            }
        } catch (Exception e) {
            log.error(e.getMessage());
        }
    }

    /**
     * 创建临时节点
     *
     * @param key
     * @param value
     */
    public void createEphemeral(String key, String value) {
        try {
            if (exists(key)) {
                remove(key);
            }
            client.create().creatingParentsIfNeeded().withMode(CreateMode.EPHEMERAL)
                    .forPath(key, value.getBytes(Charsets.UTF_8));
        } catch (Exception e) {
            log.error(e.getMessage());
        }
    }

    /**
     * 更新节点数据
     *
     * @param key
     * @param value
     */
    public void update(final String key, final String value) {
        if (!exists(key)) {
            log.error("{}节点已存在", key);
            return;
        }
        try {
            if (StringUtils.isNotBlank(value)) {
                getClient().inTransaction().check().forPath(key).and().setData()
                        .forPath(key, value.getBytes(Charsets.UTF_8)).and().commit();
            }
        } catch (Exception e) {
            log.error(e.getMessage());
        }
    }

    /**
     * 移除节点
     *
     * @param key
     */
    public void remove(String key) {
        try {
            getClient().delete().deletingChildrenIfNeeded().forPath(key);
        } catch (Exception e) {
            log.error(e.getMessage());
        }
    }

    /**
     * 获取注册中心当前时间
     *
     * @param key
     * @return
     */
    public long getCurrentTime(String key) {
        long currentTime = 0L;
        try {
            persist(key);
            currentTime = getClient().checkExists().forPath(key).getMtime();
        } catch (Exception e) {
            log.error(e.getMessage());
        }
        return currentTime;
    }

    public void increment(String key, int count, int sleepBetweenRetries) {
        try {
            distributedAtomicInteger = new DistributedAtomicInteger(getClient(), key, new RetryNTimes(count, sleepBetweenRetries));
            distributedAtomicInteger.increment();
        } catch (Exception e) {
            log.error(e.getMessage());
        }
    }

    /**
     * 获取分布式原子变量值
     *
     * @param key
     * @param count
     * @param sleepBetweenRetries
     * @return
     */
    public AtomicValue<Integer> getAtomicValue(String key, int count, int sleepBetweenRetries) {
        try {
            distributedAtomicInteger = new DistributedAtomicInteger(getClient(), key, new RetryNTimes(count, sleepBetweenRetries));
            return distributedAtomicInteger.get();
        } catch (Exception e) {
            log.error(e.getMessage());
        }
        return null;
    }

    /**
     * 添加缓存
     *
     * @param path
     */
    public void addCache(String path) {
        TreeCache cache = new TreeCache(getClient(), path);
        try {
            cache.start();
        } catch (Exception e) {
            log.error(e.getMessage());
        }
        caches.put(path + GlobalConstant.Symbol.SLASH, cache);
    }

    /**
     * 清除缓存
     *
     * @param path
     */
    public void evictCache(String path) {
        TreeCache cache = caches.remove(path + GlobalConstant.Symbol.SLASH);
        if (Objects.nonNull(cache)) {
            cache.close();
        }
    }


    public CuratorFramework getClient() {
        if (Objects.isNull(client)) {
            throw new RegException("注册中心未初始化，请调用ZookeeperRegistryCenter的init()方法进行初始化");
        }
        return client;
    }

    /**
     * 获取缓存
     *
     * @param key
     * @return
     */
    private TreeCache findTreeCache(String key) {
        for (String path : caches.keySet()) {
            if (StringUtils.startsWith(key, path)) {
                return caches.get(path);
            }
        }
        return null;
    }
}
